package adkr

import (
	"ebbflow/pkg/crypto"
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"
)

// ADKROldConfig represents the configuration for ADKR old protocol
type ADKROldConfig struct {
	SID        string        `json:"sid"`
	PID        int           `json:"pid"`
	COld       []int         `json:"c_old"`
	CNew       []int         `json:"c_new"`
	F          int           `json:"f"`
	K          int           `json:"k"`
	G          *crypto.G1    `json:"g"`
	Type       string        `json:"type"` // 's' for standard, 'b' for BN254
	SPK2s      []crypto.G1   `json:"spk2s"`
	SSK2       crypto.ZR     `json:"ssk2"`
	EPKs       []crypto.G1   `json:"epks"`
	ESK        crypto.ZR     `json:"esk"`
	THSKOld    crypto.ZR     `json:"thsk_old"`
	THPKOld    crypto.G1     `json:"thpk_old"`
	THPKsOld   []crypto.G1   `json:"thpks_old"`
}

// ADKROldOutput represents the output of ADKR old protocol
type ADKROldOutput struct {
	NewTHPK     crypto.G1     `json:"new_thpk"`
	NewTHPKs    []crypto.G1   `json:"new_thpks"`
	NewTHSK     crypto.ZR     `json:"new_thsk"`
	ConfigChain [][]int       `json:"config_chain"`
	Sigma       crypto.G1     `json:"sigma"`
	Sigma2      crypto.G1     `json:"sigma2"`
	Digest2     []byte        `json:"digest2"`
}

// ADKROldMessage represents messages in ADKR old protocol
type ADKROldMessage struct {
	Tag     string `json:"tag"`
	Round   int    `json:"round"`
	Payload []byte `json:"payload"`
}

// ADKROldTag represents different message types in ADKR old protocol
type ADKROldTag string

const (
	ADKR_ACSS_OLD   ADKROldTag = "ADKR_ACSS"
	ADKR_MVBA_OLD   ADKROldTag = "ADKR_MVBA"
	ADKR_CONFIG_OLD ADKROldTag = "ADKR_CONFIG"
	ADKR_PERM_OLD   ADKROldTag = "ADKR_PERM"
)

// ADKROldProtocol implements the old ADKR protocol
type ADKROldProtocol struct {
	config           ADKROldConfig
	msgInCh          <-chan Message
	msgOutCh         chan<- Message
	outputCh         chan<- ADKROldOutput
	logger           *log.Logger
	mu               sync.RWMutex
	acssOutputs      map[int]ACSSOutput
	completedDealers map[int]struct{}
	roundResults     []ACSSOutput
	selectedOutputs  []int
}

// NewADKROldProtocol creates a new ADKR old protocol instance
func NewADKROldProtocol(
	config ADKROldConfig,
	msgInCh <-chan Message,
	msgOutCh chan<- Message,
	outputCh chan<- ADKROldOutput,
	logger *log.Logger,
) *ADKROldProtocol {
	return &ADKROldProtocol{
		config:           config,
		msgInCh:          msgInCh,
		msgOutCh:         msgOutCh,
		outputCh:         outputCh,
		logger:           logger,
		acssOutputs:      make(map[int]ACSSOutput),
		completedDealers: make(map[int]struct{}),
		roundResults:     make([]ACSSOutput, 0),
		selectedOutputs:  make([]int, 0),
	}
}

// Run executes the ADKR old protocol
func (aop *ADKROldProtocol) Run() error {
	aop.logger.Printf("ADKROld-SID:%s PID:%d Starting ADKR old protocol", aop.config.SID, aop.config.PID)
	
	// Start message processing
	go aop.processMessages()
	
	// Run the main protocol logic
	if err := aop.runMainProtocol(); err != nil {
		return fmt.Errorf("main protocol failed: %w", err)
	}
	
	// Generate final output
	finalOutput := aop.generateFinalOutput()
	aop.outputCh <- finalOutput
	
	aop.logger.Printf("ADKROld-SID:%s PID:%d ADKR old protocol completed", aop.config.SID, aop.config.PID)
	return nil
}

// runMainProtocol runs the main protocol logic
func (aop *ADKROldProtocol) runMainProtocol() error {
	aop.logger.Printf("ADKROld-SID:%s PID:%d Running main protocol logic", aop.config.SID, aop.config.PID)
	
	// Phase 1: Run ACSS instances
	if err := aop.runACSSInstances(); err != nil {
		return fmt.Errorf("ACSS instances failed: %w", err)
	}
	
	// Phase 2: Run MVBA to select outputs
	if err := aop.runMVBA(); err != nil {
		return fmt.Errorf("MVBA failed: %w", err)
	}
	
	// Phase 3: Compute new key shares
	if err := aop.computeNewKeyShares(); err != nil {
		return fmt.Errorf("key share computation failed: %w", err)
	}
	
	// Phase 4: Generate and broadcast configuration
	if err := aop.generateAndBroadcastConfig(); err != nil {
		return fmt.Errorf("config generation failed: %w", err)
	}
	
	return nil
}

// runACSSInstances runs ACSS instances for all parties
func (aop *ADKROldProtocol) runACSSInstances() error {
	aop.logger.Printf("ADKROld-SID:%s PID:%d Running ACSS instances", aop.config.SID, aop.config.PID)
	
	N_o := len(aop.config.COld)
	N_n := len(aop.config.CNew)
	
	// Create channels for ACSS communication
	acssInChs := make([]chan Message, N_o)
	acssOutputChs := make([]chan ACSSOutput, N_o)
	acssInputChs := make([]chan [2]crypto.ZR, N_o)
	
	for i := 0; i < N_o; i++ {
		acssInChs[i] = make(chan Message, 100)
		acssOutputChs[i] = make(chan ACSSOutput, 1)
		acssInputChs[i] = make(chan [2]crypto.ZR, 1)
	}
	
	// Start ACSS instances
	var wg sync.WaitGroup
	for i := 0; i < N_o; i++ {
		wg.Add(1)
		go func(dealerID int) {
			defer wg.Done()
			
			var inputCh <-chan [2]crypto.ZR
			if aop.config.COld[dealerID] == aop.config.PID {
				inputCh = acssInputChs[dealerID]
			}
			
			CompleteSecretSharing(
				aop.config.SID+"ACSS"+fmt.Sprintf("%d", dealerID),
				aop.config.PID, N_o, aop.config.F, aop.config.COld, N_n, aop.config.F, aop.config.CNew,
				aop.config.G, aop.config.G, dealerID,
				aop.config.EPKs, aop.config.ESK,
				inputCh, acssOutputChs[dealerID], acssInChs[dealerID], aop.msgOutCh,
			)
		}(i)
	}
	
	// Provide input for our own ACSS instance
	if aop.isInOldConfig() {
		ourIndex := aop.getIndexInOldConfig()
		secretA := crypto.ZRFromInt(int64(123 + aop.config.PID))
		secretB := crypto.ZRFromInt(int64(456 + aop.config.PID))
		acssInputChs[ourIndex] <- [2]crypto.ZR{secretA, secretB}
	}
	
	// Collect ACSS outputs
	for i := 0; i < N_o; i++ {
		select {
		case output := <-acssOutputChs[i]:
			aop.mu.Lock()
			aop.acssOutputs[i] = output
			aop.completedDealers[i] = struct{}{}
			aop.mu.Unlock()
			aop.logger.Printf("ADKROld-SID:%s PID:%d Received ACSS output from dealer %d", 
				aop.config.SID, aop.config.PID, i)
		case <-time.After(10 * time.Second):
			aop.logger.Printf("ADKROld-SID:%s PID:%d Timeout waiting for ACSS output from dealer %d", 
				aop.config.SID, aop.config.PID, i)
		}
	}
	
	wg.Wait()
	return nil
}

// runMVBA runs MVBA to select which ACSS outputs to use
func (aop *ADKROldProtocol) runMVBA() error {
	aop.logger.Printf("ADKROld-SID:%s PID:%d Running MVBA", aop.config.SID, aop.config.PID)
	
	N_o := len(aop.config.COld)
	
	// Create MVBA channels
	mvbaInputCh := make(chan []byte, 1)
	mvbaOutputCh := make(chan []byte, 1)
	mvbaInCh := make(chan Message, 100)
	localSetCh := make(chan int, N_o)
	
	// Define predicate for MVBA
	predicate := func(value []byte, localSet map[int]struct{}) bool {
		var proposalSet []int
		if err := json.Unmarshal(value, &proposalSet); err != nil {
			return false
		}
		
		// Check if proposal set is valid
		if len(proposalSet) < aop.config.F+1 {
			return false
		}
		
		// Check if all proposed dealers are in local set
		for _, dealerID := range proposalSet {
			if _, ok := localSet[dealerID]; !ok {
				return false
			}
		}
		
		return true
	}
	
	// Start MVBA
    // 优先使用外部注册的阈值 BLS 配置；未注册则回退本地演示生成
    pks, thPK := GetThresholdBLSPKs()
    var mySK crypto.ZR
    if GetMyBLSSK() != nil { mySK = *GetMyBLSSK() }
    if len(pks) != N_o || thPK == nil || GetMyBLSSK() == nil {
        genPK, pksGen, skShares := crypto.GenerateThresholdKeys(N_o, aop.config.F+1)
        pks = pksGen
        mySK = skShares[aop.config.PID]
        thPK = genPK
    }
    go TThresholdMVBA(
        aop.config.SID+"MVBA",
        aop.config.PID, N_o, aop.config.F, aop.config.COld,
        aop.config.G, mvbaInputCh, mvbaOutputCh, mvbaInCh, aop.msgOutCh,
        nil, localSetCh, predicate, pks, mySK, thPK,
    )
	
	// Provide local set
	aop.mu.RLock()
	for dealerID := range aop.completedDealers {
		localSetCh <- dealerID
	}
	aop.mu.RUnlock()
	close(localSetCh)
	
	// Propose a set of dealers
	proposalSet := aop.selectProposalSet()
	proposalBytes, _ := json.Marshal(proposalSet)
	mvbaInputCh <- proposalBytes
	
	// Wait for MVBA decision
	select {
	case decisionBytes := <-mvbaOutputCh:
		var selectedDealers []int
		if err := json.Unmarshal(decisionBytes, &selectedDealers); err != nil {
			return fmt.Errorf("failed to unmarshal MVBA decision: %w", err)
		}
		
		aop.mu.Lock()
		aop.selectedOutputs = selectedDealers
		aop.mu.Unlock()
		
		aop.logger.Printf("ADKROld-SID:%s PID:%d MVBA decided on dealers: %v", 
			aop.config.SID, aop.config.PID, selectedDealers)
		
	case <-time.After(30 * time.Second):
		return fmt.Errorf("MVBA timed out")
	}
	
	return nil
}

// computeNewKeyShares computes new key shares from selected ACSS outputs
func (aop *ADKROldProtocol) computeNewKeyShares() error {
	aop.logger.Printf("ADKROld-SID:%s PID:%d Computing new key shares", aop.config.SID, aop.config.PID)
	
	aop.mu.RLock()
	selectedOutputs := aop.selectedOutputs
	acssOutputs := aop.acssOutputs
	aop.mu.RUnlock()
	
	// This is a simplified implementation
	// In a full implementation, this would involve complex polynomial interpolation
	// and threshold key reconstruction
	
	// For now, we'll use the first selected output as the new key
	if len(selectedOutputs) > 0 {
		if output, ok := acssOutputs[selectedOutputs[0]]; ok {
			aop.mu.Lock()
			aop.roundResults = append(aop.roundResults, output)
			aop.mu.Unlock()
		}
	}
	
	return nil
}

// generateAndBroadcastConfig generates and broadcasts the new configuration
func (aop *ADKROldProtocol) generateAndBroadcastConfig() error {
	aop.logger.Printf("ADKROld-SID:%s PID:%d Generating and broadcasting config", aop.config.SID, aop.config.PID)
	
	// Compute new threshold public key
	newTHPK := aop.computeNewThresholdPublicKey()
	
	// Generate signature
	digest := crypto.HashToG1([]byte(fmt.Sprintf("%v", newTHPK)))
	digest2 := crypto.HashToG1([]byte(fmt.Sprintf("%v%v", newTHPK, aop.config.CNew)))
	
	// Sign the digest
	sigma := crypto.MulG1(&digest, &aop.config.THSKOld)
	sigma2 := crypto.MulG1(&digest2, &aop.config.THSKOld)
	
	// Convert digest2 to bytes
	digest2Bytes := digest2.G1Affine.Bytes()
	
	// Create configuration message
	configMsg := ADKROldMessage{
		Tag:   string(ADKR_CONFIG_OLD),
		Round: 0,
		Payload: []byte(fmt.Sprintf("config_%s", aop.config.SID)),
	}
	
	// Broadcast configuration
	N_o := len(aop.config.COld)
	for i := 0; i < N_o; i++ {
		payloadBytes, _ := json.Marshal(configMsg)
		aop.msgOutCh <- Message{
			Tag:     BroadcastTag(ADKR_CONFIG_OLD),
			Round:   0,
			To:      i,
			Payload: payloadBytes,
		}
	}
	
	// Store the result
	aop.mu.Lock()
	defer aop.mu.Unlock()
	
	// This would normally be stored in the protocol state
	_ = ADKROldOutput{
		NewTHPK:     newTHPK,
		NewTHPKs:    aop.config.THPKsOld, // Simplified
		ConfigChain: [][]int{aop.config.CNew},
		Sigma:       sigma,
		Sigma2:      sigma2,
		Digest2:     digest2Bytes[:],
	}
	
	return nil
}

// processMessages processes incoming messages
func (aop *ADKROldProtocol) processMessages() {
	for msg := range aop.msgInCh {
		// Process different message types
		switch msg.Tag {
		case BroadcastTag(ADKR_ACSS_OLD):
			aop.logger.Printf("ADKROld-SID:%s PID:%d Received ACSS message", aop.config.SID, aop.config.PID)
		case BroadcastTag(ADKR_MVBA_OLD):
			aop.logger.Printf("ADKROld-SID:%s PID:%d Received MVBA message", aop.config.SID, aop.config.PID)
		case BroadcastTag(ADKR_CONFIG_OLD):
			aop.logger.Printf("ADKROld-SID:%s PID:%d Received CONFIG message", aop.config.SID, aop.config.PID)
		default:
			aop.logger.Printf("ADKROld-SID:%s PID:%d Received unknown message with tag %s", 
				aop.config.SID, aop.config.PID, msg.Tag)
		}
	}
}

// Helper functions

func (aop *ADKROldProtocol) isInOldConfig() bool {
	for _, pid := range aop.config.COld {
		if pid == aop.config.PID {
			return true
		}
	}
	return false
}

func (aop *ADKROldProtocol) getIndexInOldConfig() int {
	for i, pid := range aop.config.COld {
		if pid == aop.config.PID {
			return i
		}
	}
	return -1
}

func (aop *ADKROldProtocol) selectProposalSet() []int {
	aop.mu.RLock()
	defer aop.mu.RUnlock()
	
	// Select first f+1 dealers that have completed ACSS
	var selected []int
	for dealerID := range aop.completedDealers {
		if len(selected) >= aop.config.F+1 {
			break
		}
		selected = append(selected, dealerID)
	}
	return selected
}

func (aop *ADKROldProtocol) computeNewThresholdPublicKey() crypto.G1 {
	// Simplified implementation - in practice this would involve proper interpolation
	var result crypto.G1
	result.SetInfinity()
	
	// Use existing threshold public key as base
	result = aop.config.THPKOld
	
	return result
}

func (aop *ADKROldProtocol) generateFinalOutput() ADKROldOutput {
	aop.mu.RLock()
	defer aop.mu.RUnlock()
	
	return ADKROldOutput{
		NewTHPK:     aop.config.THPKOld,
		NewTHPKs:    aop.config.THPKsOld,
		ConfigChain: [][]int{aop.config.CNew},
	}
}

// ADKROldC is the main function that runs the old ADKR protocol
// This corresponds to the Python function ADKR_old_c
func ADKROldC(
	sid string, pid int, cOld, cNew []int, f, k int, g *crypto.G1, typeStr string,
	sPK2s []crypto.G1, sSK2 crypto.ZR, ePKs []crypto.G1, eSK crypto.ZR,
	thskOld crypto.ZR, thpkOld crypto.G1, thpksOld []crypto.G1,
	msgInCh <-chan Message, msgOutCh chan<- Message,
	logger *log.Logger,
) {
	if logger == nil {
		logger = log.Default()
	}
	
	logger.Printf("ADKROldC-SID:%s PID:%d Starting old ADKR protocol", sid, pid)
	
	// Create configuration
	config := ADKROldConfig{
		SID:      sid,
		PID:      pid,
		COld:     cOld,
		CNew:     cNew,
		F:        f,
		K:        k,
		G:        g,
		Type:     typeStr,
		SPK2s:    sPK2s,
		SSK2:     sSK2,
		EPKs:     ePKs,
		ESK:      eSK,
		THSKOld:  thskOld,
		THPKOld:  thpkOld,
		THPKsOld: thpksOld,
	}
	
	// Create output channel
	outputCh := make(chan ADKROldOutput, 1)
	
	// Create protocol instance
	protocol := NewADKROldProtocol(
		config,
		msgInCh,
		msgOutCh,
		outputCh,
		logger,
	)
	
	// Run the protocol
	if err := protocol.Run(); err != nil {
		logger.Printf("ADKROldC-SID:%s PID:%d Protocol failed: %v", sid, pid, err)
		return
	}
	
	// Wait for output
	select {
	case output := <-outputCh:
		logger.Printf("ADKROldC-SID:%s PID:%d Protocol completed successfully", sid, pid)
		logger.Printf("ADKROldC-SID:%s PID:%d New threshold public key: %v", sid, pid, output.NewTHPK)
	case <-time.After(60 * time.Second):
		logger.Printf("ADKROldC-SID:%s PID:%d Protocol timed out waiting for output", sid, pid)
	}
}
